from __future__ import print_function

import numpy as np
import os
from skimage.filter import denoise_tv_chambolle

import cv2
from keras.preprocessing.image import ImageDataGenerator
from keras.utils.generic_utils import Progbar

from utils import image_utils
from utils import data_utils


def pre_train(X, y=None, size=128):
    # if os.path.isfile('pp.npy'):
    #     X = np.load('pp.npy')
    #     return X, y

    # resize
    # X_resize = np.ndarray((X.shape[0], X.shape[1], size, size), dtype=np.uint8)
    # for i in range(X.shape[0]):
    #     for j in range(0, X.shape[1]):
    #         X_resize[i, j] = image_utils.crop_resize(X[i, j], size)
    # X = X_resize

    # # denoise
    for i in range(X.shape[0]):
        for j in range(0, X.shape[1]):
            X[i, j] = cv2.bilateralFilter(X[i, j], 5, 32, 32)

    X_resize = np.ndarray((X.shape[0], X.shape[1], 128, 128))
    for i in range(X.shape[0]):
        for j in range(0, X.shape[1]):
            X_resize[i, j] = image_utils.crop_resize(X[i, j], 128)
    X = X_resize    # progbar = Progbar(X.shape[0])  # progress bar for pre-processing status tracking
    #
    # for i in range(X.shape[0]):
    #     for j in range(X.shape[1]):
    #         X[i, j] = denoise_tv_chambolle(X[i, j], weight=0.1, multichannel=False)
    #     progbar.add(1)

    if y is not None:
        pass
    return X, y


def per_epoch(X):

    X_epoch = np.copy(X)
    X_epoch = image_utils.rotation_augmentation(X_epoch, 15)
    X_epoch = image_utils.shift_augmentation(X_epoch, 0.1, 0.1)

    print(X.shape)
    return X_epoch